# coding: utf-8
"""
@Time    : 2024/8/28 10:15
@Author  : Y.H LEE
"""

import copy
import itertools

from utils.tools import *
from sys_params import *
from grid_trainer import GridTrainer
import os

'''
grid search
    假设某些超参数之间的影响是独立的
'''


class GridSearch:

    def __init__(self, h_params, s_params, env):
        self.metric_dict = {"DISTANCE": 1e99}
        self.env = env

        self.exp_params = [(k, v) for k, v in s_params.items() if v != 'None']
        self.all_params = [p for p in list(h_params.keys())]

        self.trainer = GridTrainer()
        self.h_params = h_params

        self.best_params = {}
        self.best_results = {metric: value for metric, value in self.metric_dict.items() if
                             metric in h_params['METRIC']}

    def train_with_hyperparams_grid_search(self, save_model=False):

        count = 0
        hp_list = list(itertools.product(*[p[1] for p in self.exp_params]))  # *list用于解压列表, 取笛卡尔积
        best = self.best_results[self.h_params['METRIC'][0]]

        for i, var_p in enumerate(hp_list):
            grid_id = f'{i + 1}'
            for p_name, value in zip(self.exp_params, var_p):
                self.h_params[p_name[0]] = value
                # grid_id += f'{p_name[0]}-{value}_'

            self.h_params['train_logs_save_dir'] = train_logs_save_dir + 'grid_search/' + self.h_params['ALGORITHM']
            if not os.path.exists(self.h_params['train_logs_save_dir']):
                os.makedirs(self.h_params['train_logs_save_dir'])
            self.h_params['train_logs'] = '/' + grid_id + '.txt'

            count += 1

            print(f'exp: {count}/{len(hp_list)}')
            self.trainer.reset_res_list()
            dist = self.trainer(self.h_params, self.metric_dict, self.env, save_model)

            if dist < best:
                best = dist

                self.best_params = copy.deepcopy(self.h_params)
                self.best_results[self.h_params['METRIC'][0]] = best

        self.save_best(name='grid_search/grid_search_res')

    def train_with_hyperparams(self, save_model=False):
        pass

    def save_best(self, name):
        with open(train_logs_save_dir + name + '.txt', 'w', encoding='utf-8') as fw:
            fw.write('Best hyper_params are: \n')
            for k, v in self.best_params.items():
                fw.write(f'{k} : {v}\n')
            fw.write('\n')
            fw.write('Best results are: \n')
            for k, v in self.best_results.items():
                fw.write(f'{k} : {v}\n')


if __name__ == '__main__':
    pass
